{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GADTSyntax #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE OverloadedStrings #-}

module Language.Skirth.Forth where

import Bound ((>>>=), Scope, abstract1)
import Control.Monad (ap)
import Control.Monad.State.Strict (State, execState, modify)
import Data.Text (Text)

-- Untyped lambda calculus, with Word primops
data LC a where
  WordLit :: {-# UNPACK #-} !Word -> LC a
  WordAdd :: LC a
  WordMul :: LC a
  Var :: !a -> LC a
  App :: LC a -> LC a -> LC a
  Lam :: Scope () LC a -> LC a
  Let :: Scope () LC a -> Scope () LC a -> LC a
  deriving stock (Functor)

instance Applicative LC where
  pure = Var
  (<*>) = ap

instance Monad LC where
  WordLit w >>= _ = WordLit w
  WordAdd >>= _ = WordAdd
  WordMul >>= _ = WordMul
  Var v >>= f = f v
  App l r >>= f = App (l >>= f) (r >>= f)
  Lam e >>= f = Lam (e >>>= f)
  Let b e >>= f = Let (b >>>= f) (e >>>= f)

-- helpers with lambda constructor gore

lam :: (Eq a) => a -> LC a -> LC a
lam v = Lam . abstract1 v

let1 :: (Eq a) => a -> LC a -> LC a -> LC a
let1 bind def rest = Let (abstract1 bind def) (abstract1 bind rest)

letMany :: (Eq a, Foldable f) => LC a -> f (a, LC a) -> LC a
letMany = foldr (\(def, bind) acc -> let1 def bind acc)

-- Stack helpers

emptyStack :: LC Text
emptyStack = lam "a" (lam "b" (Var "b"))

pop :: LC Text
pop = lam "p" (App (Var "p") (lam "x" (lam "y" (Var "y"))))

push :: LC Text
push = lam "x" (lam "y" (lam "z" (App (App (Var "z") (Var "x")) (Var "y"))))

peek :: LC Text
peek = lam "p" (App (Var "p") (lam "x" (lam "y" (Var "x"))))

-- top of stack, second on stack, rest of stack
withTop2 :: (Text -> Text -> Text -> LC Text) -> LC Text
withTop2 f =
  lam
    "stack"
    ( letMany
        ( letMany
            (f "tos" "sos" "rros")
            [ ("sos", App peek (Var "ros")),
              ("rros", App pop (Var "ros"))
            ]
        )
        [ ("tos", App peek (Var "stack")),
          ("ros", App pop (Var "stack"))
        ]
    )

-- Actual operations

swapLC :: LC Text
swapLC =
  withTop2
    ( \tos sos ros ->
        App
          (App push (Var sos))
          (App (App push (Var tos)) (Var ros))
    )

addLC :: LC Text
addLC =
  withTop2
    ( \tos sos ros ->
        App
          (App push (App (App WordAdd (Var tos)) (Var sos)))
          (Var ros)
    )

mulLC :: LC Text
mulLC =
  withTop2
    ( \tos sos ros ->
        App
          (App push (App (App WordMul (Var tos)) (Var sos)))
          (Var ros)
    )

litLC :: Word -> LC Text
litLC w = lam "stack" (App (App push (WordLit w)) (Var "stack"))

dropLC :: LC Text
dropLC = pop

-- Compiler monad

newtype ForthM a = ForthM (State (LC Text) a)
  deriving newtype (Functor, Applicative, Monad)

-- What the users get to use

lit :: Word -> ForthM ()
lit w = ForthM . modify $ App (litLC w)

add :: ForthM ()
add = ForthM . modify $ App addLC

mul :: ForthM ()
mul = ForthM . modify $ App mulLC

swap :: ForthM ()
swap = ForthM . modify $ App swapLC

drop :: ForthM ()
drop = ForthM . modify $ App dropLC

compile :: ForthM () -> LC Text
compile (ForthM comp) = execState comp emptyStack
